package cipher

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"go-chaos/code"
	"go-chaos/file"
	"os"
	"os/user"
	"path/filepath"
)

type GenerateKeyHandle func(string, string)

type Sign interface {
	Sign(data []byte) ([]byte, error)
	VerySign(data, signData []byte) bool
}

type RsaPub struct {
	PubKey []byte // pki
	Msg    []byte
}

type RsaPrv struct {
	PrvKey []byte // pkcs1
	Cipher []byte
}

type RsaCipher struct {
	Pub  *RsaPub
	Prv  *RsaPrv
	Hash crypto.Hash // sign
}

// rsa public key
func (pub *RsaPub) PublicKey() *rsa.PublicKey {
	key, err := x509.ParsePKIXPublicKey(pub.PubKey)
	if err != nil {
		panic(err)
	}
	return key.(*rsa.PublicKey)
}

// rsa private key
func (prv *RsaPrv) PrivateKey() *rsa.PrivateKey {
	if key, err := x509.ParsePKCS1PrivateKey(prv.PrvKey); err == nil {
		return key
	}
	return nil
}

// rsa encrypt
func (r *RsaCipher) Encrypt() ([]byte, error) {
	pubKey := r.Pub.PublicKey()
	if pubKey == nil {
		return nil, CError{
			detail: "parse public key error",
		}
	}

	bs, err := rsa.EncryptPKCS1v15(rand.Reader, pubKey, r.Pub.Msg)
	if err != nil {
		return nil, err
	}
	return bs, nil
}

// rsa decrypt
func (r *RsaCipher) Decrypt() ([]byte, error) {
	prv := r.Prv.PrivateKey()
	if prv == nil {
		return nil, CError{
			detail: "parse private key error",
		}
	}

	bs, err := rsa.DecryptPKCS1v15(rand.Reader, prv, r.Prv.Cipher)
	if err != nil {
		return nil, err
	}
	return bs, nil
}

// rsa sign
func (r *RsaCipher) Sign(data []byte) ([]byte, error) {
	prvKey := r.Prv.PrivateKey()

	h := Hash{
		t:   r.Hash,
		src: data,
	}

	sign, err := rsa.SignPKCS1v15(rand.Reader, prvKey, r.Hash, h.Hash())
	if err != nil {
		return nil, CError{
			detail: "rsa sign had some problem",
		}
	}
	return sign, nil
}

// rsa very sign
func (r *RsaCipher) VerySign(data, signData []byte) bool {
	pubKey := r.Pub.PublicKey()
	h := Hash{
		t:   r.Hash,
		src: data,
	}
	err := rsa.VerifyPKCS1v15(pubKey, r.Hash, h.Hash(), signData)
	if err == nil {
		return true
	}
	panic(err)
}

// rsa generate public and private key
func GenerateKey(bits int, dstDir string, handle GenerateKeyHandle) error {
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return err
	}

	bs := x509.MarshalPKCS1PrivateKey(privateKey)
	prvBase64 := code.Base64Encode(bs)
	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: bs,
	}

	if dstDir == "" {
		usr, err := user.Current()
		if err != nil {
			return err
		}
		dstDir = usr.HomeDir
	} else if !file.IsDir(dstDir) && !file.PathExist(dstDir) {
		if err = os.MkdirAll(dstDir, os.ModePerm); err != nil {
			return err
		}
	}

	prvf, _ := os.Create(dstDir + string(filepath.Separator) + "private.pem")
	_ = pem.Encode(prvf, block)
	defer func() {
		_ = prvf.Close()
	}()

	// PKI比较常用
	bs, _ = x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	block = &pem.Block{
		Type:  "PUB KEY",
		Bytes: bs,
	}
	pubf, _ := os.Create(dstDir + string(filepath.Separator) + "public.pem")
	_ = pem.Encode(pubf, block)
	defer func() {
		_ = pubf.Close()
	}()

	pubBase64 := code.Base64Encode(bs)
	if handle != nil {
		handle(prvBase64, pubBase64)
	}
	return nil
}
